{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 优先级经验回放\n",
    "***离线算法总是跟经验池离不开关系***\n",
    "\n",
    "但是一般的随机抽样不能够将有用的经验反复学习，所以有了PER-Prioritized Experience Replay\n",
    "\n",
    "1. 如果算法足够好，没必要用PER.\n",
    "2. 如果奖励函数足够dense，也没有必要用PER.\n",
    "3. 因为用了PER速度大大降低，性能不一定会提高。\n",
    "4. 如果真实场景交互，CPU资源足够，稀疏奖励，可以试试。\n",
    "\n",
    "## 回顾DQN中的ER-Experience Replay\n",
    "1. 保持队列的结构，先进先出，当长度到达定值后开始踢人\n",
    "2. 从整个Buffer中均匀分布随机选择\n",
    "3. 然而，在DQN数以千计的样本，也就是 [s, a, r, s’] 中，一定有一些样本可以帮助DQN更快的收敛；或者说，存在一些样本，如果我们先学了这些样本，后面的就不用再学了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "class SumTree():\n",
    "    def __init__(self,max_mem) -> None:\n",
    "        self.capacity=max_mem\n",
    "        self.data_ptr=0\n",
    "        # 对于capacity个数据节点，我们需要使用capacity-1个父节点来形成一个树形结构\n",
    "        # 那么一共需要 2*capacity-1个节点，定义tree[:self.capacity-1]为父节点，后面的是叶子节点\n",
    "        self.tree=np.zeros(2*self.capacity-1)\n",
    "        #存储数据的节点，与tree[self.capacity-1:] 一一对应\n",
    "        self.data=np.zeros(self.capacity,dtype=object)\n",
    "    \n",
    "    def add(self,priority:float,data:object):\n",
    "        # 叶子节点的下标\n",
    "        tree_idx = self.data_ptr+self.capacity-1\n",
    "        # 存储数据\n",
    "        self.data[self.data_ptr]=data\n",
    "        # 更新父节点的权重\n",
    "        self.update_parent_weights(tree_idx,priority)\n",
    "        # 如果达到了存储的上限，则踢掉刚刚进来的数据\n",
    "        self.data_ptr+=1\n",
    "        if self.data_ptr>=self.capacity:\n",
    "            self.data_ptr=0\n",
    "\n",
    "    def update_parent_weights(self,tree_idx:int,priority:float):\n",
    "        # 获得修改的值\n",
    "        change_value = priority - self.tree[tree_idx]\n",
    "        # 直接修改叶子节点的值\n",
    "        self.tree[tree_idx]=priority\n",
    "        # 跟这个叶子节点有关的父亲节点 全部要更新\n",
    "        while tree_idx != 0:\n",
    "            tree_idx=(tree_idx-1)//2\n",
    "            self.tree[tree_idx]+=change_value\n",
    "\n",
    "    def get_leaf(self,v:float):\n",
    "        \"\"\"\n",
    "        根据采样到的值v来取叶子节点的数据\n",
    "        \"\"\"\n",
    "        parent_idx=0\n",
    "        while True:\n",
    "            cl_idx=2*parent_idx+1\n",
    "            cr_idx=cl_idx+1\n",
    "            # 超出界限了就结束搜索\n",
    "            if cl_idx>=len(self.tree):\n",
    "                leaf_idx=parent_idx\n",
    "                break\n",
    "            else:\n",
    "                if v<=self.tree[cl_idx]:\n",
    "                    parent_idx=cl_idx\n",
    "                else:\n",
    "                    v=v-self.tree[cl_idx]\n",
    "                    parent_idx=cr_idx\n",
    "        data_idx=leaf_idx-(self.capacity-1)\n",
    "        # return leaf_idx,and the priority of the idx, and the data\n",
    "        return leaf_idx,self.tree[leaf_idx],self.data[data_idx]\n",
    "\n",
    "    @property\n",
    "    def total_priority(self):\n",
    "        return self.tree[0]\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 定义PER类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProritizedExperienceReplay():\n",
    "    def __init__(self,\n",
    "    max_size:int,\n",
    "    batch_size:int,\n",
    "    clipped_abs_error:float,\n",
    "    epsilon:float,\n",
    "    alpha:float,\n",
    "    beta:float,\n",
    "    beta_increment:float,\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Params\n",
    "        ------\n",
    "        alpha:trade-off factor 控制采样在uniform和greedy的偏好,0代表均匀随机，1代表完全按贪婪算法\n",
    "        epsilon: 避免除数为0的一个很小的数\n",
    "        beta: 决定抵消PER对于收敛结果的影响，1代表完全抵消，这样就与ER没有区别了\n",
    "        beta_increment: 每次采样对beta采取的增量\n",
    "        clipped_abs_error: abs_error的上限\n",
    "        \"\"\"\n",
    "\n",
    "        #存储参数\n",
    "        self.clipped_abs_error = clipped_abs_error\n",
    "        self.batch_size = batch_size\n",
    "        self.beta_increment=beta_increment\n",
    "        self.epsilon=epsilon\n",
    "        self.alpha=alpha\n",
    "        self.beta=beta\n",
    "        # 定义SUMTREE\n",
    "        self.tree=SumTree(max_size)\n",
    "\n",
    "    def store_transition(self,transition:object):\n",
    "        # 索引后面的叶子节点\n",
    "        max_priority=np.max(self.tree.tree[-self.tree.capacity:])\n",
    "        if max_priority==0:\n",
    "            max_priority=self.clipped_abs_error\n",
    "            #存进来的时候，调成最大的优先级，保证每个样本都要过一遍\n",
    "        self.tree.add(max_priority,transition)\n",
    "    \n",
    "    def sample(self):\n",
    "        \"\"\"\n",
    "        Returns\n",
    "        -------\n",
    "        Tree index : the indexs of the tree\n",
    "        Data : the data of the nodes\n",
    "        ISWeights : the normalized Priority of the nodes\n",
    "        \"\"\"\n",
    "        b_idx=np.empty((self.batch_size),dtype=np.int32)\n",
    "        b_memory=np.empty((self.batch_size,self.tree.data[0].size))\n",
    "        ISWeights = np.empty((self.batch_size, 1))\n",
    "        #分成batch_size 个段\n",
    "        priority_segments= self.tree.total_priority/self.batch_size\n",
    "        # 递增的beta\n",
    "        self.beta=np.min([1.,self.beta+self.beta_increment])\n",
    "        # 取得最小倍率，这是normalization的过程\n",
    "        min_prob = np.min(self.tree.tree[-self.tree.capacity])/self.tree.total_priority\n",
    "\n",
    "        for i in range(self.batch_size):\n",
    "            # 将优先值分为batch_size片，然后从每片中均匀采样\n",
    "            a,b = priority_segments*i,priority_segments*(i+1)\n",
    "            sample_val=np.random.uniform(a,b)\n",
    "            idx,p,data = self.tree.get_leaf(sample_val)\n",
    "            #取得优先值的占比\n",
    "            prob=p/self.tree.total_priority\n",
    "            ISWeights[i,0]=np.power(prob/min_prob,-self.beta)\n",
    "            b_idx[i],b_memory[i,:]=idx,data\n",
    "        # 返回抽样值\n",
    "        return b_idx,b_memory,ISWeights\n",
    "\n",
    "    def batch_update(self,tree_idx:np.ndarray,abs_errors:np.ndarray):\n",
    "        # + epsilon避免除以0\n",
    "        abs_errors+=self.epsilon\n",
    "        clipped_errors=np.minimum(abs_errors,self.clipped_abs_error)\n",
    "        ps = np.power(clipped_errors,self.alpha)\n",
    "        for ti,p in zip(tree_idx,ps):\n",
    "            self.tree.update_parent_weights(ti,p)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TEST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id:0->propotion:6.71 vs target weights 6.67%\n",
      "id:1->propotion:13.33 vs target weights 13.33%\n",
      "id:2->propotion:19.96 vs target weights 20.00%\n",
      "id:3->propotion:26.90 vs target weights 26.67%\n",
      "id:4->propotion:33.09 vs target weights 33.33%\n"
     ]
    }
   ],
   "source": [
    "mem = ProritizedExperienceReplay(5,1,100.0,1e-5,1,1,0.001)\n",
    "for i in range(5):\n",
    "\n",
    "    mem.store_transition((np.array([i+1.,i+2.])))\n",
    "for i in range(5):\n",
    "    mem.batch_update(np.array([4,5,6,7,8]),np.array([1,2,3,4,5],dtype=np.float32))\n",
    "\n",
    "sum_sample = 100000\n",
    "index_count = np.zeros((5))\n",
    "for i in range(sum_sample):\n",
    "    id,_,_=mem.sample()\n",
    "    index_count[id-4]+=1\n",
    "\n",
    "for i in range(5):\n",
    "    print(\n",
    "        f\"id:{i}->propotion:{index_count[i]/sum(index_count)*100 :.2f} vs target weights {(i+1)/15 *100 :.2f}%\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 问题\n",
    "1. 如果奖励不够稀疏，它还不如随机采样，随机采样的时间复杂度基本上是O（1），而他的采样复杂度和更新复杂度都比随机高\n",
    "2. 消耗大量CPU资源\n",
    "3. 抽样不均匀可能会导致过拟合"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d7d6d563c462dbcb2ee090b6970b42fc638e8eb4a4e6decc12c92a9df43cdf15"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 64-bit ('base': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
